import tempfile
from subprocess import check_call, call, PIPE
from forest.registry import register_tool

######## Execution contexts - use with the 'with' statement ##############
@register_tool('image.DMDev')
class DMDev(object):
    def __init__(self, dev, size, name=None):
        (self.dev, self.size, self.name) = (dev, size, name)
    def __enter__(self):
        self.mapperdev = dm_attach(self.dev, self.size, self.name)
        return self.mapperdev
    def __exit__(self, exc_type, exc_value, traceback):
        dm_detach(self.mapperdev)


def dm_attach(dev, size, name=None):
    '''Attach a devicemapper device to the given device, with the given size.
    If name is None, a random name will be chosen. Returns the device name.
    raises CalledProcessError if dmsetup fails.'''
    if name is None:
        name = tempfile.mktemp(prefix="lorax.imgutils.", dir="")
    check_call(["dmsetup", "create", name, "--table",
                "0 %i linear %s 0" % (size/512, dev)],
                stdout=PIPE, stderr=PIPE)
    return name

@register_tool('image.dm_detach')
def dm_detach(dev):
    '''Detach the named devicemapper device. Returns False if dmsetup fails.'''
    dev = dev.replace("/dev/mapper/", "") # strip prefix, if it's there
    return call(["dmsetup", "remove", dev], stdout=PIPE, stderr=PIPE)

